{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jF7E8k7dD-IL"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/drive/1COw1BDnON0hUdCuugzbAFwbgEGXHgZ-f?usp=sharing\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "\n",
    "# Cache-Augmented Generation (CAG) Explained\n",
    "\n",
    "This notebook provides an explanation of Cache-Augmented Generation (CAG), its advantages, limitations, and a guide on how to set up and run the provided code for testing CAG."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pzApD2M6D-IM"
   },
   "source": [
    "## What is Cache-Augmented Generation (CAG)?\n",
    "\n",
    "Retrieval-Augmented Generation (RAG) is a powerful approach that enhances language models by integrating external knowledge. However, RAG faces challenges like retrieval latency, retrieval errors, and system complexity.\n",
    "\n",
    "**Cache-Augmented Generation (CAG)** is an alternative paradigm designed to bypass real-time retrieval. It leverages the extended context windows of modern Large Language Models (LLMs) by preloading all relevant resources into the model's context and caching its runtime parameters (specifically, the Key-Value (KV) cache). During inference, the preloaded KV-cache allows the model to generate responses directly, eliminating the need for dynamic retrieval.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1ZPQSLssO1En"
   },
   "source": [
    "### Advantages of CAG\n",
    "\n",
    "- **Reduced Latency**: Eliminates real-time retrieval, leading to faster inference.\n",
    "- **Improved Reliability**: Minimizes retrieval errors while maintaining context relevance.\n",
    "- **Simplified Design**: Offers a streamlined, retrieval-free alternative to RAG, capable of achieving comparable or superior results with lower complexity."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iJBW52SFO49A"
   },
   "source": [
    "### Limitations of CAG\n",
    "\n",
    "- **Limited Context Size Window of LLM**: CAG requires the entire knowledge source to fit within the context window, making it less suitable for tasks involving extremely large datasets.\n",
    "- **Context being too large**: The performance of LLMs may degrade with very long contexts."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2ITbx0-xPu4M"
   },
   "source": [
    "## **Architectural Overview of CAG and RAG**\n",
    "<br>\n",
    "\n",
    "![Difference between RAG and CAG](https://raw.githubusercontent.com/hhhuang/CAG/refs/heads/main/overview.png)\n",
    "\n",
    "In CAG, the query is appended to the KV Cache and then LLM generates the answer and after that KV Cache is reset to original by truncating query.\n",
    "So, that the KV Cache size does not keep on increasing, if it would have happened then it would have consumed the whole LLM Context Window.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZcS4dDhuD-IM"
   },
   "source": [
    "## Installation\n",
    "\n",
    "To get started, first ensure you have the necessary dependencies installed. You can do this by running the following command in your terminal:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rV59yR-kD-IN"
   },
   "outputs": [],
   "source": [
    "!git clone https://github.com/hhhuang/CAG.git\n",
    "%cd CAG\n",
    "%pwd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vpBneU9uD-IO"
   },
   "outputs": [],
   "source": [
    "!pip install -r ./requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "itb5r0RhGd0F"
   },
   "outputs": [],
   "source": [
    "!pip uninstall torch torchvision torchaudio --yes\n",
    "!pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "F72vv6UPD-IO"
   },
   "source": [
    "## Preparation\n",
    "\n",
    "Before running the CAG experiments, you need to perform a couple of preparation steps:\n",
    "\n",
    "### 1. Download Datasets\n",
    "Download the required `squad` and `hotpotqa` datasets by running the provided curl script:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fZC2sgzMD-IO"
   },
   "outputs": [],
   "source": [
    "!sh ./downloads.sh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Fs-ZSqINTU4u"
   },
   "source": [
    "Datasets used:\n",
    "- SQuAD: Dataset which focuses on precise, context-aware answers withing single passages\n",
    "- HotPotQA: Dataset which focuses on multihop reasoning questions across multiple documents."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CdhxG8x5D-IO"
   },
   "source": [
    "### 2. Configure Environment Variables\n",
    "Insert your `HF_TOKEN` for Hugging Face authentication. If `HF_TOKEN` is not found, the script will raise an error."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6G4kwmngOWgi"
   },
   "source": [
    "To add your HF_TOKEN key to Colab secrets:\n",
    "\n",
    "1. Click on the \"🔑 Secrets\" tab in the left sidebar of your Colab notebook.\n",
    "2. Click the \"+\" button to add a new secret.\n",
    "3. In the \"Name\" field, enter `HF_TOKEN`.\n",
    "4. In the \"Value\" field, paste your HF Token.\n",
    "5. Make sure the \"Notebook access\" toggle is enabled for this secret.\n",
    "6. Click \"Done\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "kexfaH-rIGIX"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from google.colab import userdata\n",
    "os.environ['HF_TOKEN'] = userdata.get('HF_TOKEN')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "I5dlai77D-IP"
   },
   "source": [
    "Create a HF_TOKEN with Write Access, create a new token at https://huggingface.co/settings/tokens\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "v_Nv9RUFT4HC"
   },
   "source": [
    "## Setup\n",
    "\n",
    "Model used here for testing CAG and RAG is **Llama 3.2-1B Model** which has a context size of 128k tokens\n",
    "\n",
    "In RAG, we comapre the results of CAG with 2 metrics:\n",
    "  - **BM25 (Sparse Retrieval):**  BM25,a sparse retrieval algorithm, ranks documents based on term frequency inverse document frequency (TF-IDF) and document length normalization\n",
    "  - **Dense Retrieval System (OpenAI Indexes):**  For a query q, dense retrieval selects the top-k passages p that semantically align with the query,\n",
    " offering improved contextual understanding compared to sparse methods."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WXOd-lygD-IP"
   },
   "source": [
    "## Running CAG Experiments with `kvcache.py`\n",
    "\n",
    "The `kvcache.py` script is designed for running CAG experiments. It preprocesses knowledge into a KV cache and then uses this cache for generating responses to questions, evaluating performance based on semantic similarity and generation time."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HBe48mD4D-IP"
   },
   "source": [
    "### Key Concepts from `kvcache.py`\n",
    "\n",
    "- **KV Cache (Key-Value Cache)**: In transformer models, the KV cache stores the 'keys' and 'values' from previous token computations. This prevents re-computation for each new token during sequential decoding, significantly speeding up inference.\n",
    "- **`preprocess_knowledge` function**: This function takes a knowledge prompt, tokenizes it, and runs it through the model to generate the initial KV cache (lines 80-99 in `kvcache.py`). This precomputed cache represents the 'augmented knowledge' in CAG.\n",
    "- **`write_kv_cache` and `read_kv_cache` functions**: These functions handle saving and loading the precomputed KV cache to/from disk (lines 102-126 in `kvcache.py`). This allows for persistent storage and reuse of the knowledge cache.\n",
    "- **`generate` function**: This function performs greedy decoding using the pre-existing `past_key_values` (KV cache) to generate new tokens (lines 40-77 in `kvcache.py`). In the CAG context, this `past_key_values` is the knowledge cache.\n",
    "- **`clean_up` function**: This function truncates the KV cache to its original length, which is important when reusing the same knowledge cache for multiple queries to prevent it from growing indefinitely (lines 114-118 in `kvcache.py`)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xOD6oA8FD-IP"
   },
   "source": [
    "### `kvcache.py` Parameters\n",
    "\n",
    "Here are the important parameters you can use with `kvcache.py`:\n",
    "\n",
    "- `--kvcache`: Specifies the KV cache method. Use `\"file\"` to read/write cache from/to a file.\n",
    "- `--dataset`: The dataset to use for experiments. Options include `\"hotpotqa-train\"`, `\"squad-train\"`, `\"kis\"`, `\"kis_sample\"`, `\"squad-dev\"`, `\"hotpotqa-dev\"`, `\"hotpotqa-test\"`\n",
    "- `--similarity`: The similarity metric for evaluation. Currently, `\"bertscore\"` is supported.\n",
    "- `--modelname`: The name of the Hugging Face model to use (e.g., `\"meta-llama/Llama-3.1-8B-Instruct\"`).\n",
    "- `--maxKnowledge`:  Integer. Selects how many documents from the dataset to use as knowledge.\n",
    "- `--maxParagraph`:  Integer. Limits the number of paragraphs per knowledge document (default: 100).\n",
    "- `--maxQuestion`:  Integer. Specifies the maximum number of questions to test.\n",
    "- `--randomSeed`:  Integer. A random seed number for reproducibility.\n",
    "- `--output`: String. The filepath to save the results of the experiment.\n",
    "- `--usePrompt`: (Flag) Add this parameter if you do NOT want to use the CAG knowledge cache acceleration and instead include the knowledge directly in the prompt (disables caching)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "02UPJQ1PD-IP"
   },
   "source": [
    "### Notes on Parameters\n",
    "\n",
    "#### `--maxKnowledge` parameter notice:\n",
    "Approximate Tokens count corresponding to knowledge document size of `\"squad-train\"` and `\"hotpotqa-train\"` dataset.\n",
    "\n",
    "- **`datasets=(\"squad-train\")`**\n",
    "  - when k = 3, tokens = 21,000\n",
    "  - when k = 4, tokens = 32,000\n",
    "  - when k = 7, tokens = 50,000\n",
    "\n",
    "- **`datasets=(\"hotpotqa-train\")`**\n",
    "  - all k = 7405 article, tokens = 10,038,084\n",
    "  - when k = 1, tokens = 1,400\n",
    "  - when k = 16, tokens = 22,400\n",
    "  - when k = 24, tokens = 33,667\n",
    "  - when k = 32, tokens = 44,800\n",
    "  - when k = 48, tokens = 64,000\n",
    "  - when k = 64, tokens = 85,000\n",
    "  - when k = 80, tokens = 106,000\n",
    "\n",
    "#### `--maxQuestion` parameter notice:\n",
    "- When using `\"squad-train\"` dataset, 1 knowledge has an average of 150 questions.\n",
    "- When using `\"hotpotqa-train\"` dataset, 1 knowledge has 1 question.\n",
    "\n",
    "> **TIP:** Since 1 document in `\"hotpoqa-train\"` dataset has only 1 question, it may not satisfy large-scale evaluation. Multiple evaluations could be a relatively better approach."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "f_t2m28XD-IP"
   },
   "source": [
    "### Example Usage - CAG\n",
    "\n",
    "To run a CAG experiment, you can use a command similar to this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ZrN9S0SOD-IP"
   },
   "outputs": [],
   "source": [
    "!python ./kvcache.py --kvcache file --dataset \"squad-train\" --similarity bertscore \\\n",
    "    --maxKnowledge 1 --maxParagraph 1 --maxQuestion 1 \\\n",
    "    --modelname \"meta-llama/Llama-3.2-8B-Instruct\" --randomSeed 0 \\\n",
    "    --output \"./result_kvcache.txt\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "toKoXYsdWRIM"
   },
   "source": [
    "## Running RAG Experiments with `rag.py`\n",
    "\n",
    "For comparison purposes, this repository also includes a traditional RAG implementation in `rag.py`. This allows you to benchmark CAG performance against conventional retrieval-augmented generation approaches."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "L-hfMh3QWSSW"
   },
   "source": [
    "### Key Differences: RAG vs CAG\n",
    "\n",
    "| Aspect | RAG | CAG |\n",
    "|--------|-----|-----|\n",
    "| **Retrieval** | Dynamic, on-demand | Pre-loaded, cached |\n",
    "| **Latency** | Higher (retrieval overhead) | Lower (no retrieval step) |\n",
    "| **Knowledge Size** | Can handle large datasets | Limited by context window |\n",
    "| **Complexity** | Higher system complexity | Simpler, streamlined |\n",
    "| **Memory Usage** | Lower during inference | Higher (cached knowledge) |\n",
    "| **Retrieval Errors** | Possible | Eliminated |\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_v8VJ_2IWfrL"
   },
   "source": [
    "### `rag.py` Parameters\n",
    "\n",
    "Here are the important parameters you can use with `rag.py`:\n",
    "\n",
    "- `--index`: The indexing method for retrieval. Options: `\"openai\"` or `\"bm25\"`\n",
    "- `--dataset`: The dataset to use for experiments. Options: `\"hotpotqa-train\"` or `\"squad-train\"`\n",
    "- `--similarity`: The similarity metric for evaluation. Currently, `\"bertscore\"` is supported.\n",
    "- `--maxKnowledge`: (Optional) Integer. Selects how many documents from the dataset to use as knowledge.\n",
    "- `--maxParagraph`: (Optional) Integer. Limits the number of paragraphs per knowledge document (default: 100).\n",
    "- `--maxQuestion`: (Optional) Integer. Specifies the maximum number of questions to test.\n",
    "- `--topk`: Integer. The number of top similar documents to retrieve for each query.\n",
    "- `--modelname`: The name of the Hugging Face model to use (e.g., `\"meta-llama/Llama-3.1-8B-Instruct\"`).\n",
    "- `--randomSeed`: (Optional) Integer. A random seed number for reproducibility.\n",
    "- `--output`: String. The filepath to save the results of the experiment."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "J2vjwfiPWgTO"
   },
   "source": [
    "### Example Usage - RAG\n",
    "\n",
    "To run a RAG experiment for comparison with CAG, you can use a command similar to this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BWekqNutWpl4"
   },
   "outputs": [],
   "source": [
    "!python ./rag.py --index \"bm25\" --dataset \"hotpotqa-train\" --similarity bertscore \\\n",
    "    --maxKnowledge 80 --maxParagraph 100 --maxQuestion 80 --topk 3 \\\n",
    "    --modelname \"meta-llama/Llama-3.2-8B-Instruct\" --randomSeed 0 \\\n",
    "    --output \"./rag_results.txt\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "F4L4_e1leVWp"
   },
   "source": [
    "### **NOTE**: You can download the ouptut results from  \"Files\" tab in the left sidebar of your Colab notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rzxqj6JCXfcL"
   },
   "source": [
    "## Mind Map of CAG Process\n",
    "![Full Process.jpg]()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
